#! /usr/bin/env python
# Calling: ./aufgabe3.py abab | dot -Tpng -o tmp.png; eog tmp.png

import sys
import argparse
import re
from sets import Set
from collections import defaultdict


def main(argv):
    parser = argparse.ArgumentParser(description='Compute the A+ tree for the sequence seq')
    parser.add_argument('seq', help='sequence to process')
    args = parser.parse_args()

    s_tree(args.seq)


def getLongestCommonFrefix(l):
    l.sort(key=len)

    lcp = ''
#    print l
#    print l[0]
    for i in range(1, len(l[0]) + 1):
        regex = '^' + l[0][0:i] + '.*'
        rmatch = re.compile(regex)

        listLength = len(l)
        filteredList = filter(rmatch.match, l)

        if listLength == len(filteredList):
            lcp = l[0][0:i]
        else:
            break

    return lcp


def computeNodes(suffixList, parentNode, parentLabel):
    print '\t' + parentNode + ' -> ' + parentLabel + 'X [label=X]'
    # Computing list of chars in the suffix list
    charList = list()

    s = Set()
    for i in suffixList:
        s = s | Set(i)

    for item in s:
        charList.append(item)

    # Computing R-Sets
    hitList = defaultdict(list)
    for i in charList:
        regex = '^' + i + '.*'
        rmatch = re.compile(regex)
        hitList[i] = filter(rmatch.match, suffixList)

    # Computing the longest common prefixes for each R-group
    for k in hitList:
        if len(hitList[k]) == 0:
            break
        elif len(hitList[k]) == 1:
            n = hitList[k]
            print '\t' + parentNode + ' -> ' + parentLabel + n[0] + 'X' + ' [label=' + n[0] + 'X]'
        else:
            node = getLongestCommonFrefix(hitList[k])
            print '\t' + parentNode + ' -> ' + parentLabel + node + ' [label=' + node + ']'

            prefixLength = len(node)

            # Cutting preffix
            hitList[k] = [x[prefixLength:] for x in hitList[k]]

            # Removing empty elements
            hitList[k] = [x for x in hitList[k] if x]

#            sys.stdout.write("calling computeNodes(")
#            print hitList[k]
            computeNodes(hitList[k], node, node)


def s_tree(seq):

    # Computing list of suffixes
    tmp = seq
    suffixList = list()
    while tmp:
        suffixList.append(tmp)
        tmp = tmp[1:]

    print 'digraph suffixtree_for_' + seq + '{'

    computeNodes(suffixList, 'root', '')

    print '}'

if __name__ == "__main__":
    main(sys.argv)
